/* 
 * Solves the Panfilov model using an explicit numerical scheme.
 * Based on code orginally provided by Xing Cai, Simula Research Laboratory
 * 
 * Modified and  restructured by Scott B. Baden, UCSD
 */

#include <assert.h>
#include <stdlib.h>
#include <iostream>
#include <iomanip>
#include <string.h>
#include <math.h>
#include <sys/time.h>
#include <mpi.h> 
using namespace std;


// Utilities
// 

// Timer
// Make successive calls and take a difference to get the elapsed time.
static const double kMicro = 1.0e-6;
double getTime()
{
    struct timeval TV;
    struct timezone TZ;

    const int RC = gettimeofday(&TV, &TZ);
    if(RC == -1) {
            cerr << "ERROR: Bad call to gettimeofday" << endl;
            return(-1);
    }

    return( ((double)TV.tv_sec) + kMicro * ((double)TV.tv_usec) );

}  // end getTime()

// Allocate a 2D array
double **alloc2D(int m,int n){
   double **E;
   int nx=n+1, ny=m+1;
   E = (double**)malloc(sizeof(double*)*ny + sizeof(double)*nx*ny);
   assert(E);
   int j;
   for(j=0;j<ny;j++) E[j] = (double*)(E+ny) + j*nx;
   return(E);
}
    
// Reports statistics about the computation
// These values should not vary (except to within roundoff)
// when we use different numbers of  processes to solve the problem
 double stats(double **E, int m, int n, double *_mx){
     double mx = -1;
     double l2norm = 0;
     int i, j;
     for (j=1; j<=m+1; j++)
       for (i=1; i<=n+1; i++) {
	   l2norm += E[j][i]*E[j][i];
	   if (E[j][i] > mx)
	       mx = E[j][i];
      }
     *_mx = mx;
     l2norm /= (double) ((m+1)*(n+1));
     l2norm = sqrt(l2norm);
     return l2norm;
 }

void gather(double** from,double** to, double** temp,int cutsWide,int cutsHigh,int w,int h){
    int size = cutsWide*cutsHigh;
    //if horizontal strips
    if (cutsWide == 1){
        MPI_Gather(from[1],h*(w+2),MPI_DOUBLE,to[1],h*(w+2),MPI_DOUBLE,0,MPI_COMM_WORLD);
    }
    //if vertical strips
    else {
        MPI_Send(from[1],h*(w+2),MPI_DOUBLE,0,0,MPI_COMM_WORLD);
        int rank;
        MPI_Comm_rank(MPI_COMM_WORLD,&rank);
        if (rank == 0){
            int source;
            int start;
            int end;
            for (int i = 0;i < size;i++){
                MPI_Status status;
                MPI_Recv(temp[1],h*(w+2),MPI_DOUBLE,MPI_ANY_SOURCE,0,MPI_COMM_WORLD,&status);
                source = status.MPI_SOURCE;
                start = w*source+1;
                end = start + w;
                for (int j = 1;j < h;j++)
                    for (int i = start;i < end;i++)
                        to[j][i] = temp[j][i-start+1];
            }
        }
    }
}

// External functions
extern "C" {
    void splot(double **E, double T, int niter, int m, int n);
}
void cmdLine(int argc, char *argv[], double& T, int& n, int& px, int& py, int& do_stats, double& plot_freq, int& no_comm);


// Main program
int main (int argc, char** argv)
{
 /*
  *  Solution arrays
  *   E is the "Excitation" variable, a voltage
  *   R is the "Recovery" variable
  *   E_prev is the Excitation variable for the previous timestep,
  *      and is used in time integration
  *   E_final is used for plotting.
  */
 double **E, **R, **E_prev, **E_final;

 // Various constants - these definitions shouldn't change
 const double a=0.1, b=0.1, kk=8.0, M1= 0.07, M2=0.3, epsilon=0.01, d=5e-5;

 double T=1000.0;
 int m=200,n=200;
 int do_stats = 0;
 double plot_freq = 0.0;
 int px = 1, py = 1;
 int no_comm = 0;

 cmdLine( argc, argv, T, n,px, py, do_stats,  plot_freq, no_comm);
 m = n;


 MPI_Init(&argc,&argv);

 int size;
 int rank;
 MPI_Comm_size(MPI_COMM_WORLD,&size);
 MPI_Comm_rank(MPI_COMM_WORLD,&rank);
 if (size < px*py){
     printf("not enough processors, need %d, but have %d",px*py,size);
     return 0;
 }
 size = px*py;
 if (size == 1)
     no_comm = 1;
if (rank >= size){//any extra processors terminate early.
    MPI_Finalize();
    return 1;
}

 E_final = alloc2D(m+2,n+2);

 m = (m+1)/py - 1;
 n = (n+1)/px - 1;

 // Allocate contiguous memory for solution arrays
 // The computational box is defined on [1:m+1,1:n+1]
 // We pad the arrays in order to facilitate differencing on the 
 // boundaries of the computation box
 E = alloc2D(m+2,n+2);
 E_prev = alloc2D(m+2,n+2);
 R = alloc2D(m+2,n+2);
 int i,j;
 // Initialization
 for (j=1; j<=m+1; j++)
   for (i=1; i<=n+1; i++)
     E_prev[j][i] = R[j][i] = 0.0;
 if (size == 1){//single processor
     for (j=1;j<=m+1;j++)
         for (i=n/2+2;i<=n+1;i++)
             E_prev[j][i] = 1.0;
     for (j=m/2+2;j<=m+1;j++)
         for (i=1;i<=n+1;i++)
             R[j][i] = 1.0;
 }
 else if (px == 1){//horizontal slices
     for (j=1; j<=m+1; j++)
         for (i=n/2+2; i<=n+1; i++)
             E_prev[j][i] = 1.0;
     if (rank >= size/2){//top half
         for (j=1;j <= m+1;j++)
             for (i = 1; i <= n+1;i++)
                 R[j][i] = 1.0;
     }
 } else if (py == 1){//Split up vertically
     for (j = m/2+2;j <= m+1;j++)
         for (i = 1;i <= n+1;i++)
             R[j][i] = 1.0;
     if (rank >= size/2){//top half
         for (j = 1;j <= m+1;j++)
             for (i = 1;i <= n+1;i++)
                 R[j][i] = 1.0;
     }
 } else {//can't do squares.......
      printf("Can't do no squares, need bigger paycheck for that");
      return 0;
 }
 double dx = 1.0/n;

 // For time integration, these values shouldn't change 
 double rp= kk*(b+1)*(b+1)/4;
 double dte=(dx*dx)/(d*4+((dx*dx))*(rp+kk));
 double dtr=1/(epsilon+((M1/M2)*rp));
 double dt = (dte<dtr) ? 0.95*dte : 0.95*dtr;
 double alpha = d*dt/(dx*dx);
 if (rank == 0){
     cout << "dt= " << dt << ", T = " << T << endl;
     cout << "processor geometry: " << px << " x " << py << endl;
     if (no_comm)
         cout << " *** Communication DISABLED" << endl;
     cout << endl;
 }
 // Start the timer
 double t0 = -getTime();


 // Simulated time is different from the integer timestep number
 // Simulated time
 double t = 0.0;
 // Integer timestep number
 int niter=0;
 double* leftGhost;
 double* rightGhost;
 if (py == 1){
     if(rank!=0)
         leftGhost = (double*)malloc(sizeof(double)*(m+3));
     if(rank!=size-1)
         rightGhost = (double*)malloc(sizeof(double)*(m+3));
 }
 
 int numReq;
 if (rank == 0 || rank == size - 1)
     numReq = 1;
 else
     numReq = 2;
 MPI_Request* requests = (MPI_Request*)malloc(sizeof(MPI_Request)*numReq);
 MPI_Status* statuses = (MPI_Status*)malloc(sizeof(MPI_Status)*numReq);
 while (t<T) {
   t += dt;
   niter++;

   /* 
    * Copy data from boundary of the computational box 
    * to the padding region, set up for differencing
    * on the boundary of the computational box
    *
    */

//if left edge is a boundary
if (no_comm || px == 1 || rank == 0)
    #pragma ivdep
        for (j=1; j<=m+1; j++) 
            E_prev[j][0] = E_prev[j][2];
else {//if left edge is ghost cells
    for (j = 0;j <= m+2;j++)
        leftGhost[j] = E_prev[j][1];
    MPI_Send(leftGhost,m+3,MPI_DOUBLE,rank-1,0,MPI_COMM_WORLD);
    MPI_Irecv(leftGhost,m+3,MPI_DOUBLE,rank-1,0,MPI_COMM_WORLD,requests);
}

//if right edge is a boundary
if (no_comm|| px == 1 || rank == size-1)
    #pragma ivdep
        for (j=1; j<=m+1; j++) 
            E_prev[j][n+2] = E_prev[j][n];
else{//if right edge is ghost cells
    for (j = 0;j <= m+2;j++)
        rightGhost[j] = E_prev[j][n+1];
    MPI_Send(rightGhost,m+3,MPI_DOUBLE,rank+1,0,MPI_COMM_WORLD);
    if (rank == 0)
        MPI_Irecv(rightGhost,m+3,MPI_DOUBLE,rank+1,0,MPI_COMM_WORLD,requests);
    else
        MPI_Irecv(rightGhost,m+3,MPI_DOUBLE,rank+1,0,MPI_COMM_WORLD,requests+1);
}

//if low edge is boundary
if (no_comm || py == 1 || rank == 0)
    #pragma ivdep
        for (i=1; i<=n+1; i++) 
            E_prev[0][i] = E_prev[2][i];
else{//if low edge is ghost cells
    MPI_Isend(E_prev[1],n+3,MPI_DOUBLE,rank-1,0,MPI_COMM_WORLD,requests);
    MPI_Irecv(E_prev[0],n+3,MPI_DOUBLE,rank-1,0,MPI_COMM_WORLD,requests);
}

//if high edge is boundary
if (no_comm || py == 1 || rank == size - 1)
    #pragma ivdep
        for (i=1; i<=n+1; i++) 
            E_prev[m+2][i] = E_prev[m][i];
else{//if high edge is ghost cells
    MPI_Send(E_prev[m+1],n+3,MPI_DOUBLE,rank+1,0,MPI_COMM_WORLD);
    if (rank == 0)
        MPI_Irecv(E_prev[m+2],n+3,MPI_DOUBLE,rank+1,0,MPI_COMM_WORLD,requests);
    else
        MPI_Irecv(E_prev[m+2],n+3,MPI_DOUBLE,rank+1,0,MPI_COMM_WORLD,requests+1);
}

if (!no_comm){
    MPI_Waitall(numReq,requests,statuses);
//    for (i = 0;i < numReq;i++)
//        MPI_Status stat = statuses[i];
//        if (stat != MPI_SUCCESS){
//            printf("Error in communication of ghost cells in processor %d",rank);
//            return 1;
//        }
}

//Copy from right/left Ghosts into actual ghost cells if necessary
if (py == 1 && rank != 0)
    for (j = 1;j <= m+1;j++)
        E_prev[j][0] = leftGhost[j];

if (py == 1 && rank != size - 1)
    for (j = 1;j <= m+1;j++)
        E_prev[j][n+2] = rightGhost[j];

   // Solve for the excitation, the PDE
   for (j=1; j<=m+1; j++){
#pragma ivdep
     for (i=1; i<=n+1; i++) {
	E[j][i] = E_prev[j][i]+alpha*(E_prev[j][i+1]+E_prev[j][i-1]-4*E_prev[j][i]+E_prev[j+1][i]+E_prev[j-1][i]);
     }
   }

   /* 
    * Solve the ODE, advancing excitation and recovery to the
    *     next timtestep
    */
   for (j=1; j<=m+1; j++){
     double *RR = &R[j][1];
     double *EE = &E[j][1];
     for (i=1; i<=n+1; i++, EE++,RR++)
	EE[0] += -dt*(kk*EE[0]*(EE[0]-a)*(EE[0]-1)+EE[0]*RR[0]);
   }

   for (j=1; j<=m+1; j++){
     double *RR = &R[j][1];
     double *EE = &E[j][1];
     for (i=1; i<=n+1; i++, EE++, RR++)
	RR[0] += dt*(epsilon+M1* RR[0]/( EE[0]+M2))*(-RR[0]-kk*EE[0]*(EE[0]-b-1));
   }

    if(!no_comm)
        MPI_Barrier(MPI_COMM_WORLD);

//   k = (int)(t/100);
//   int k = (int)(t/10);
   if (do_stats && rank==0){
//       if ((t-k*100)<dt) 
        if ((niter %  1000) == 0 ){
	 cout << "iteration " << niter << ", t = " << t << endl;
       }
   }

   if (plot_freq){
        int k = (int)(t/plot_freq);
        if ((t-k*plot_freq)<dt && no_comm)
            splot(E,t,niter,m+1,n+1);
        else if ((t-k*plot_freq)<dt && !no_comm){
            gather(E,E_final,E_prev,px,py,n+1,m+1);
            if (rank == 0)
                splot(E_final,t,niter,(m+1)*py,(n+1)*px);
        }
//       sleep(0.5);
//       scanf("%d",&resp);
   }
   double **tmp = E; E = E_prev; E_prev = tmp;
 }


 MPI_Barrier(MPI_COMM_WORLD);
 if (rank == 0){
     t0 += getTime();
     cout <<  "End at time: " << T << ",  iteration " << niter << endl;
     double mx;
     double l2norm = stats(E_prev,m,n,&mx);
     printf("Max: %13.6e, L2norm: %13.6e\n",mx,l2norm);
     cout << "Wallclock time: " << t0 << " sec." << endl;
     if (plot_freq){
         cout << "\n\nEnter any input to close the program and the plot..." << endl;
         int resp;
         scanf("%d",&resp);
     }
 }

 free (E);
 free (E_prev);
 free (R);
 MPI_Finalize();
}
